from langchain.llms import OpenAI
from langchain.prompts import PromptTemplate
from langchain.chains import LLMChain,SimpleSequentialChain

script_prompt_tpl = PromptTemplate.from_template(
    '你是一个优秀的编剧。请使用你丰富的想象力根据我定的标题编写一个故事概要。'
    '标题：{title}'
)
script_llm = OpenAI(
    api_key='hk-mtiquv1000041663a49a34520ad3294132cbce7abf1c2ef3',
    base_url="https://api.openai-hk.com/v1"
)
script_chain = LLMChain(llm=script_llm,prompt =script_prompt_tpl)

#创建广告链
adv_prompt_tpl = PromptTemplate.from_template(
    '你是一个优秀的广告写手，根据故事的概要。'
    '为我的故事写一个尽可能短但要让人有观看欲望的广告词'
    '故事概要：{story}'
)
adv_llm = OpenAI(
    api_key='hk-mtiquv1000041663a49a34520ad3294132cbce7abf1c2ef3',
    base_url="https://api.openai-hk.com/v1"
)
adv_chain = LLMChain(llm=adv_llm,prompt =adv_prompt_tpl)

chain = SimpleSequentialChain(chains = [script_chain,adv_chain],verbose=True)
print(chain.run('孙悟空大战变形金刚'))